package com.xydpeng.tech.member.framework.shardingjdbc;

import org.apache.shardingsphere.infra.exception.core.ShardingSpherePreconditions;
import org.apache.shardingsphere.infra.exception.generic.UnsupportedSQLOperationException;
import org.apache.shardingsphere.sharding.api.sharding.standard.PreciseShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.RangeShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.StandardShardingAlgorithm;
import org.apache.shardingsphere.sharding.exception.data.NullShardingValueException;

import java.util.Collection;
import java.util.HashMap;
import java.util.Objects;
import java.util.Properties;

/**
 * @author xydpeng
 * @version v1.0
 * @date: 2024/8/5
 * @description:
 */
public class CustomStandardShardingAlgorithm implements StandardShardingAlgorithm<Comparable<?>> {
    private static final String ALLOW_RANGE_QUERY_KEY = "allow-range-query-with-custom-sharding";
    private static final String TOTAL_TABLE_NUM_KEY = "totalTableNum";
    private static final String TABLES_IN_DATABASE_KEY = "tablesInDatabase";

    private boolean allowRangeQuery = false;

    private int totalTableNum;

    private HashMap<Integer, String> tablesInDatabase;

    @Override
    public void init(Properties props) {
        allowRangeQuery = isAllowRangeQuery(props);
        totalTableNum = getTotalTableNum(props);
        tablesInDatabase = getTablesInDatabase(props);
        if (totalTableNum != tablesInDatabase.size()) {
            throw new IllegalArgumentException("CustomStandardShardingAlgorithm props wrong: totalTableNum != tablesInDatabase.size()");
        }
    }


    @Override
    public String doSharding(Collection<String> availableTargetNames, PreciseShardingValue<Comparable<?>> shardingValue) {
        ShardingSpherePreconditions.checkNotNull(shardingValue.getValue(), NullShardingValueException::new);
        String columnName = shardingValue.getColumnName();
        Comparable<?> value = shardingValue.getValue();
        int index = Math.abs(value.hashCode()) % totalTableNum;
        return tablesInDatabase.get(index);
    }


    @Override
    public Collection<String> doSharding(Collection<String> availableTargetNames, RangeShardingValue<Comparable<?>> shardingValue) {
        ShardingSpherePreconditions.checkState(allowRangeQuery,
                () -> new UnsupportedSQLOperationException(String.format("Since the property of `%s` is false, inline sharding algorithm can not tackle with range query", ALLOW_RANGE_QUERY_KEY)));
        //如果允许范围查询，则会广播sql语句
        return availableTargetNames;
    }


    private boolean isAllowRangeQuery(final Properties props) {
        return Boolean.parseBoolean(props.getOrDefault(ALLOW_RANGE_QUERY_KEY, Boolean.FALSE.toString()).toString());
    }

    private int getTotalTableNum(Properties props) {
        Object property = props.get(TOTAL_TABLE_NUM_KEY);
        if (Objects.isNull(property)) {
            throw new IllegalArgumentException("CustomStandardShardingAlgorithm props wrong: totalTableNum is null");
        }
        return Integer.parseInt(String.valueOf(property).trim());
    }

    private HashMap<Integer, String> getTablesInDatabase(Properties props) {
        String property = props.getProperty(TABLES_IN_DATABASE_KEY);
        if (Objects.isNull(property)) {
            throw new IllegalArgumentException("CustomStandardShardingAlgorithm props wrong: totalTableNum is null");
        }
        HashMap<Integer, String> tablesInDatabase = new HashMap<>();
        String[] split = property.split(",");
        for (String s : split) {
            String[] kv = s.split("=");
            String ds = kv[0].trim();
            String[] ts = kv[1].replace("[", "").replace("]", "").trim().split(",");
            for (String t : ts) {
                if (t.contains("..")) {
                    String[] rang = t.split("\\.\\.");
                    for (int i = Integer.parseInt(rang[0].trim()); i <= Integer.parseInt(rang[1].trim()); i++) {
                        tablesInDatabase.put(i, ds);
                    }
                } else {
                    tablesInDatabase.put(Integer.parseInt(t.trim()), ds);
                }
            }
        }
        return tablesInDatabase;
    }


}
